##########################################################################################

library(data.table)
library(optparse)
library(dplyr)

##########################################################################################

option_list <- list(
    make_option(c("--gene"), type = "character") ,
    make_option(c("--pre_file"), type = "character") ,
    make_option(c("--cancer_file"), type = "character") ,
    make_option(c("--sample_info"), type = "character") ,
    make_option(c("--gtf_file"), type = "character") ,
    make_option(c("--out_path"), type = "character")
)

if(1!=1){
    
    gene <- "MUC6"
    pre_file <- "~/20220915_gastric_multiple/dna_combinePublic/maf/All_GGA.precancer.maf"
    cancer_file <- "~/20220915_gastric_multiple/dna_combinePublic/maf/All_GGA.cancer.maf"
    sample_info <- "~/20220915_gastric_multiple/dna_combinePublic/baseTable/STAD_Info.addBurden.MSI_MSS.addCNVType.tsv"
    out_path <- paste0("~/20220915_gastric_multiple/dna_combinePublic/images/lollipop" , "/" , gene , "_NMU")
    gtf_file <- "~/ref/GTF/gencode.v19.annotation.exonNum.gtf"

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

pre_file <- opt$pre_file
cancer_file <- opt$cancer_file
sample_info <- opt$sample_info
gtf_file <- opt$gtf_file
gene <- opt$gene
out_path <- opt$out_path

###########################################################################################

dir.create(out_path , recursive = T)

###########################################################################################

dat_PreCancerous <- fread(pre_file,sep = "\t" ,  quote="" ,header = T)
dat_Cancerous <- fread(cancer_file  ,sep = "\t",  quote="" ,header = T)
dat_info <- data.frame(fread(sample_info))
dat_gtf <- fread(gtf_file)

###########################################################################################

Variant_Types <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")

###########################################################################################
## 突变整理
dat_Cancerous$t_alt_count <- as.numeric(dat_Cancerous$t_alt_count)
dat_Cancerous$t_ref_count <- as.numeric(dat_Cancerous$t_ref_count)
dat_PreCancerous$t_alt_count <- as.numeric(dat_PreCancerous$t_alt_count)
dat_PreCancerous$t_ref_count <- as.numeric(dat_PreCancerous$t_ref_count)

dat_Cancerous <- subset(dat_Cancerous , t_alt_count > 0 & Hugo_Symbol == gene & Variant_Classification %in% Variant_Types )
dat_PreCancerous <- subset(dat_PreCancerous , t_alt_count > 0 & Hugo_Symbol == gene & Variant_Classification %in% Variant_Types )

if(1!=1){
    ## 存在其它转录本
    ## 使用第一个转录本，因为EGFR基因的默认转录本就在第一个
    maf <- data.frame(dat_Cancerous)
    if(nrow(dat_Cancerous) > 0){
        maf_use <- data.frame(
            gene = maf$Hugo_Symbol, 
            refseq = sapply( strsplit(sapply(strsplit(maf$Other_Transcripts , "[|]") , "[" , 1) , "_") , "[" , 2 ) ,
        	chromosome = paste0( "chr" , maf$Chromosome) , start = maf$Start_Position , 
            REF = maf$Tumor_Seq_Allele1 , ALT = maf$Tumor_Seq_Allele2 ,
        	aachange = paste0( "p." , sapply( strsplit(sapply(strsplit(maf$Other_Transcripts , "[|]") , "[" , 1) , "p[.]") , "[" , 2 )) ,
        	class = maf$Variant_Classification , 
            t_alt_count = maf$t_alt_count , t_ref_count = maf$t_ref_count ,
            VAF = maf$t_alt_count/(maf$t_alt_count + maf$t_ref_count) ,
            Tumor = maf$Tumor_Sample_Barcode
            )

        maf_use_mutiple <- maf_use
    }

    maf <- data.frame(dat_PreCancerous)
    if(nrow(dat_PreCancerous) > 0){
        maf_use <- data.frame(
            gene = maf$Hugo_Symbol, 
            refseq = sapply( strsplit(sapply(strsplit(maf$Other_Transcripts , "[|]") , "[" , 1) , "_") , "[" , 2 ) ,
            chromosome = paste0( "chr" , maf$Chromosome) , start = maf$Start_Position , 
            REF = maf$Tumor_Seq_Allele1 , ALT = maf$Tumor_Seq_Allele2 ,
            aachange = paste0( "p." , sapply( strsplit(sapply(strsplit(maf$Other_Transcripts , "[|]") , "[" , 1) , "p[.]") , "[" , 2 )) ,
            class = maf$Variant_Classification , 
            t_alt_count = maf$t_alt_count , t_ref_count = maf$t_ref_count ,
            VAF = maf$t_alt_count/(maf$t_alt_count + maf$t_ref_count) ,
            Tumor = maf$Tumor_Sample_Barcode

            )
        maf_use_single <- maf_use
    }

    if(nrow(dat_Cancerous) > 0 & nrow(dat_PreCancerous) > 0 ){
        maf_use <- rbind( maf_use_mutiple , maf_use_single )
    }else if(nrow(dat_Cancerous) > 0){
        maf_use <- maf_use_mutiple
    }else if(nrow(dat_PreCancerous) > 0){
        maf_use <- maf_use_single
    }

    maf_use$refseq <- sapply(strsplit(maf_use$refseq,"[.]") , "[" , 1)
}

## 使用软件默认转录本
maf <- data.frame(dat_Cancerous)
if(nrow(dat_Cancerous) > 0){
    maf_use <- data.frame(
        gene = maf$Hugo_Symbol, 
        refseq = maf$Annotation_Transcript ,
        chromosome = paste0( "chr" , maf$Chromosome) , start = maf$Start_Position , 
        REF = maf$Tumor_Seq_Allele1 , ALT = maf$Tumor_Seq_Allele2 ,
        aachange = maf$Protein_Change ,
        class = maf$Variant_Classification , 
        t_alt_count = maf$t_alt_count , t_ref_count = maf$t_ref_count ,
        VAF = maf$t_alt_count/(maf$t_alt_count + maf$t_ref_count) ,
        Tumor = maf$Tumor_Sample_Barcode
        )

    maf_use_mutiple <- maf_use
}

maf <- data.frame(dat_PreCancerous)
if(nrow(dat_PreCancerous) > 0){
    maf_use <- data.frame(
        gene = maf$Hugo_Symbol, 
        refseq = maf$Annotation_Transcript ,
        chromosome = paste0( "chr" , maf$Chromosome) , start = maf$Start_Position , 
        REF = maf$Tumor_Seq_Allele1 , ALT = maf$Tumor_Seq_Allele2 ,
        aachange = maf$Protein_Change ,
        class = maf$Variant_Classification , 
        t_alt_count = maf$t_alt_count , t_ref_count = maf$t_ref_count ,
        VAF = maf$t_alt_count/(maf$t_alt_count + maf$t_ref_count) ,
        Tumor = maf$Tumor_Sample_Barcode

        )
    maf_use_single <- maf_use
}

if(nrow(dat_Cancerous) > 0 & nrow(dat_PreCancerous) > 0 ){
    maf_use <- rbind( maf_use_mutiple , maf_use_single )
}else if(nrow(dat_Cancerous) > 0){
    maf_use <- maf_use_mutiple
}else if(nrow(dat_PreCancerous) > 0){
    maf_use <- maf_use_single
}

maf_use$refseq <- sapply(strsplit(maf_use$refseq,"[.]") , "[" , 1)

############################################################################################################
## 标记基线信息
maf_use <- merge( maf_use , dat_info , by = "Tumor" )

## 若不影响蛋白质改变，则用class替换其aachange
#maf_use[maf_use$aachange=="","aachange"] <- maf_use[maf_use$aachange=="","class"]
#maf_use[maf_use$aachange=="p.NA","aachange"] <- maf_use[maf_use$aachange=="p.NA","class"]
index <- which(maf_use$class=="Splice_Site")

maf_use[maf_use$class=="Splice_Site","aachange"] <- paste("Splice_Site" , 
    maf_use$chromosome[index], maf_use$start[index] , maf_use$REF[index] , maf_use$ALT[index] , 
    sep = ":")

maf_use[maf_use$class=="Splice_Site","aachange"] <- "Splice_Site"

############################################################################################################
## 总的
lollipop_variant <- maf_use
lollipop_variant_pre <- subset( lollipop_variant , Class %in% c("IM") )
lollipop_variant_can <- subset( lollipop_variant , Class %in% c("IGC" , "DGC") )

############################################################################################################
## 注释其外显子
## 其对应的外显子
maf_use$exon_pos <- ""
for(i in 1:nrow(maf_use)){
    exon_num <- dat_gtf[which(dat_gtf$transcript_id == maf_use$refseq[i] & dat_gtf$chr == maf_use$chromosome[i] & dat_gtf$start <= maf_use$start[i] & dat_gtf$end >= maf_use$start[i]  ),"exon_number"]
    exon_num <- as.numeric(unique(data.frame(exon_num)))
    maf_use$exon_pos[i] <- exon_num    
}

############################################################################################################
## 将突变分为share和Private
maf_use$vid <- paste( maf_use$chromosome , maf_use$start , maf_use$REF , maf_use$ALT , sep = ":" )

result <- c()
for( normal in unique(maf_use$Normal) ){
    tmp <- subset( maf_use , Normal == normal )

    for( mut in unique(tmp$vid) ){
        tmp2 <- subset( tmp , vid == mut )
        share <- length(grep( "IM" , unique(tmp2$Class))) > 0 & length(grep( "GC" , unique(tmp2$Class))) > 0
        if(share){
            tmp2$Share <- TRUE
        }else{
            tmp2$Share <- FALSE
        }
        result <- rbind( result , tmp2 )
    }
}

maf_use <- result

## 所有的信息
out_name <- paste0(out_path , "/" , gene , ".AllInfo.tsv" )
write.table(maf_use , out_name , row.names=F , sep ="\t" , quote = F)

## Tumor
out_name <- paste0(out_path , "/" , gene , ".PreCancerous.tsv" )
write.table(lollipop_variant_pre , out_name , row.names=F , sep ="\t" , quote = F)

out_name <- paste0(out_path , "/" , gene , ".Cancerous.tsv" )
write.table(lollipop_variant_can , out_name , row.names=F , sep ="\t" , quote = F)

############################################################################################################
## 输出画图使用，以Normal为单位避免重复
## 按照Normal分类
lollipop_variant_pre_sample <- unique(
    lollipop_variant_pre[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )

lollipop_variant_cancer_sample <- unique(
    lollipop_variant_can[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )

## Normal
#lollipop_variant_pre_sample <- subset( lollipop_variant_pre_sample , Type != "IM + IGC + DGC" )
out_name <- paste0(out_path , "/" , gene , ".PreCancerous.UniqueNormal.tsv" )
write.table(lollipop_variant_pre_sample , out_name , row.names=F , sep ="\t" , quote = F)

## Tumor
#lollipop_variant_cancer_sample <- subset( lollipop_variant_cancer_sample , Type != "IM + IGC + DGC" )
out_name <- paste0(out_path , "/" , gene , ".Cancerous.UniqueNormal.tsv" )
write.table(lollipop_variant_cancer_sample , out_name , row.names=F , sep ="\t" , quote = F)

## IGC
lollipop_variant_can_igc <- subset( lollipop_variant_can , Class=="IGC")
lollipop_variant_can_igc_sample <- unique(
    lollipop_variant_can_igc[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )
out_name <- paste0(out_path , "/" , gene , ".IGC.UniqueNormal.tsv" )
write.table(lollipop_variant_can_igc_sample , out_name , row.names=F , sep ="\t" , quote = F)

## DGC
lollipop_variant_can_dgc <- subset( lollipop_variant_can , Class=="DGC")
lollipop_variant_can_dgc_sample <- unique(
    lollipop_variant_can_dgc[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )
out_name <- paste0(out_path , "/" , gene , ".DGC.UniqueNormal.tsv" )
write.table(lollipop_variant_can_dgc_sample , out_name , row.names=F , sep ="\t" , quote = F)


## IGC
lollipop_variant_can_igc <- subset( lollipop_variant_can , Class=="IGC" & Type=="IM + IGC")
lollipop_variant_can_igc_sample <- unique(
    lollipop_variant_can_igc[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )
out_name <- paste0(out_path , "/" , gene , ".IGC.IM_IGC.UniqueNormal.tsv" )
write.table(lollipop_variant_can_igc_sample , out_name , row.names=F , sep ="\t" , quote = F)

## DGC
lollipop_variant_can_dgc <- subset( lollipop_variant_can , Class=="DGC" & Type=="IM + DGC")
lollipop_variant_can_dgc_sample <- unique(
    lollipop_variant_can_dgc[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )
out_name <- paste0(out_path , "/" , gene , ".DGC.IM_DGC.UniqueNormal.tsv" )
write.table(lollipop_variant_can_dgc_sample , out_name , row.names=F , sep ="\t" , quote = F)

## IGC
lollipop_variant_can_igc <- subset( lollipop_variant_can , Class=="IGC" & Type=="IM + IGC + DGC")
lollipop_variant_can_igc_sample <- unique(
    lollipop_variant_can_igc[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )
out_name <- paste0(out_path , "/" , gene , ".IGC.IM_IGC_DGC.UniqueNormal.tsv" )
write.table(lollipop_variant_can_igc_sample , out_name , row.names=F , sep ="\t" , quote = F)

## DGC
lollipop_variant_can_dgc <- subset( lollipop_variant_can , Class=="DGC" & Type=="IM + IGC + DGC")
lollipop_variant_can_dgc_sample <- unique(
    lollipop_variant_can_dgc[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )
out_name <- paste0(out_path , "/" , gene , ".DGC.IM_IGC_DGC.UniqueNormal.tsv" )
write.table(lollipop_variant_can_dgc_sample , out_name , row.names=F , sep ="\t" , quote = F)


############################################################################################################
## IGC的IM
lollipop_variant_pre_igc <- subset( lollipop_variant_pre_sample , Type=="IM + IGC")
out_name <- paste0(out_path , "/" , gene , ".IM_IGC.UniqueNormal.tsv" )
write.table(lollipop_variant_pre_igc , out_name , row.names=F , sep ="\t" , quote = F)

## DGC的IM
lollipop_variant_pre_dgc <- subset( lollipop_variant_pre_sample , Type=="IM + DGC")
out_name <- paste0(out_path , "/" , gene , ".IM_DGC.UniqueNormal.tsv" )
write.table(lollipop_variant_pre_dgc , out_name , row.names=F , sep ="\t" , quote = F)

## IGC_DGC的IM
lollipop_variant_pre_dgc <- subset( lollipop_variant_pre_sample , Type=="IM + IGC + DGC")
out_name <- paste0(out_path , "/" , gene , ".IM_IGC-DGC.UniqueNormal.tsv" )
write.table(lollipop_variant_pre_dgc , out_name , row.names=F , sep ="\t" , quote = F)


############################################################################################################
## 共享和私有
lollipop_variant_trunk <- subset( maf_use , Share == TRUE )
lollipop_variant_private <- subset( maf_use , Share == FALSE )

## 共享
## 癌前
lollipop_variant_trunk_pre <- subset( lollipop_variant_trunk , Class == "IM"  )
lollipop_variant_trunk_pre_sample <- unique(
    lollipop_variant_trunk_pre[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )
## 癌
lollipop_variant_trunk_cancer <- subset( lollipop_variant_trunk , Class %in% c("IGC" , "DGC")  )
lollipop_variant_trunk_cancer_sample <- unique(
    lollipop_variant_trunk_cancer[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )


## 私有
## 癌前
lollipop_variant_private_pre <- subset( lollipop_variant_private , Class == "IM"  )
lollipop_variant_private_pre_sample <- unique(
    lollipop_variant_private_pre[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )
## 癌
lollipop_variant_private_cancer <- subset( lollipop_variant_private , Class %in% c("IGC" , "DGC")  )
lollipop_variant_private_cancer_sample <- unique(
    lollipop_variant_private_cancer[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )


## Tumor
out_name <- paste0(out_path , "/" , gene , ".Cancer.Trunk.tsv" )
write.table(lollipop_variant_trunk_cancer_sample , out_name , row.names=F , sep ="\t" , quote = F)
out_name <- paste0(out_path , "/" , gene , ".Cancer.Private.tsv" )
write.table(lollipop_variant_private_cancer_sample , out_name , row.names=F , sep ="\t" , quote = F)

## Pre
out_name <- paste0(out_path , "/" , gene , ".Precancer.Trunk.tsv" )
write.table(lollipop_variant_trunk_pre_sample , out_name , row.names=F , sep ="\t" , quote = F)
out_name <- paste0(out_path , "/" , gene , ".Precancer.Private.tsv" )
write.table(lollipop_variant_private_pre_sample , out_name , row.names=F , sep ="\t" , quote = F)


############################################################################################################
## 分不同的分子亚型
cin_sample <- unique( subset( dat_info , TCGA_Class=="CIN" )$Patient )
gs_sample <- unique( subset( dat_info , TCGA_Class=="GS" )$Patient )

lollipop_variant_cin_igc <- subset( maf_use , Type == "IM + IGC" & Class == "IM" & Patient %in% cin_sample )
lollipop_variant_gs_igc <- subset( maf_use , Type == "IM + IGC" & Class == "IM" & Patient %in% gs_sample )

lollipop_variant_cin_dgc <- subset( maf_use , Type == "IM + DGC" & Class == "IM" & Patient %in% cin_sample )
lollipop_variant_gs_dgc <- subset( maf_use , Type == "IM + DGC" & Class == "IM" & Patient %in% gs_sample )

lollipop_variant_cin_igc_sample <- unique(
    lollipop_variant_cin_igc[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )
lollipop_variant_gs_igc_sample <- unique(
    lollipop_variant_gs_igc[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )
lollipop_variant_cin_dgc_sample <- unique(
    lollipop_variant_cin_dgc[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )
lollipop_variant_gs_dgc_sample <- unique(
    lollipop_variant_gs_dgc[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )

out_name <- paste0(out_path , "/" , gene , ".IM_IGC.IM.CIN.tsv" )
write.table(lollipop_variant_cin_igc_sample , out_name , row.names=F , sep ="\t" , quote = F)
out_name <- paste0(out_path , "/" , gene , ".IM_IGC.IM.GS.tsv" )
write.table(lollipop_variant_gs_igc_sample , out_name , row.names=F , sep ="\t" , quote = F)

out_name <- paste0(out_path , "/" , gene , ".IM_DGC.IM.CIN.tsv" )
write.table(lollipop_variant_cin_dgc_sample , out_name , row.names=F , sep ="\t" , quote = F)
out_name <- paste0(out_path , "/" , gene , ".IM_DGC.IM.GS.tsv" )
write.table(lollipop_variant_gs_dgc_sample , out_name , row.names=F , sep ="\t" , quote = F)

## GC中
lollipop_variant_cin_igc <- subset( maf_use , Type == "IM + IGC" & Class == "IGC" & Patient %in% cin_sample )
lollipop_variant_gs_igc <- subset( maf_use , Type == "IM + IGC" & Class == "IGC" & Patient %in% gs_sample )

lollipop_variant_cin_dgc <- subset( maf_use , Type == "IM + DGC" & Class == "DGC" & Patient %in% cin_sample )
lollipop_variant_gs_dgc <- subset( maf_use , Type == "IM + DGC" & Class == "DGC" & Patient %in% gs_sample )

lollipop_variant_cin_igc_sample <- unique(
    lollipop_variant_cin_igc[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )
lollipop_variant_gs_igc_sample <- unique(
    lollipop_variant_gs_igc[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )
lollipop_variant_cin_dgc_sample <- unique(
    lollipop_variant_cin_dgc[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )
lollipop_variant_gs_dgc_sample <- unique(
    lollipop_variant_gs_dgc[, c( "Normal" , "gene" , "refseq" , "chromosome" , "start" , "REF" , "ALT" , "aachange" , "class" , "Type" ) ]
    )

out_name <- paste0(out_path , "/" , gene , ".IM_IGC.IGC.CIN.tsv" )
write.table(lollipop_variant_cin_igc_sample , out_name , row.names=F , sep ="\t" , quote = F)
out_name <- paste0(out_path , "/" , gene , ".IM_IGC.IGC.GS.tsv" )
write.table(lollipop_variant_gs_igc_sample , out_name , row.names=F , sep ="\t" , quote = F)

out_name <- paste0(out_path , "/" , gene , ".IM_DGC.DGC.CIN.tsv" )
write.table(lollipop_variant_cin_dgc_sample , out_name , row.names=F , sep ="\t" , quote = F)
out_name <- paste0(out_path , "/" , gene , ".IM_DGC.DGC.GS.tsv" )
write.table(lollipop_variant_gs_dgc_sample , out_name , row.names=F , sep ="\t" , quote = F)